#pragma once

#include <opencv2/opencv.hpp>

#include "param.h"

// 装甲板类
struct Armor {
  cv::RotatedRect bar1; // 灯条1
  cv::RotatedRect bar2; // 灯条2
  std::string type; // 装甲板类型
  int confidence_level;

  Armor(const cv::RotatedRect& bar1, const cv::RotatedRect& bar2,
        const std::string& type = "Unknown",
        const int& confidence_level = 0)
      : bar1(bar1), bar2(bar2), type(type), confidence_level(confidence_level) {}
};

// 装甲板识别
class ArmorDetect {
private:
  bool enemy; // 敌人颜色
  cv::Mat img; // 输出图像
  cv::Mat src; // 原始图像
  
  // BGR 转灰度图
  void BGR2Gray(cv::Mat& src, cv::Mat& gray); 
  // 图像增强和二值化
  void imageEnhance(cv::Mat& gray, cv::Mat& enhance); 
  // 轮廓检测，筛选出灯条
  void findLightBar(cv::Mat& enhance, std::vector<cv::RotatedRect>& light_bars); 
  // 灯条配对，筛选出装甲板
  void matchLightBar(std::vector<cv::RotatedRect>& light_bars, std::vector<Armor>& armors); 
  // 获取装甲板图像
  void getArmorImg(const cv::Mat& src, const Armor& armor, cv::Mat& armor_img);

public:
  ArmorDetect(bool enemy = param::RED) : enemy(enemy) {}
  // 设置敌人颜色
  void setEnemy(bool enemy); 
  // 灯条检测主函数
  void detect(cv::Mat& src, cv::Mat& dst);
};